from os import close
from backtrader import indicator
from backtrader.functions import And
import pandas as pd
from pymongo import MongoClient
import datetime
import time
import math
import backtrader as bt
import backtrader.indicators as btind
import backtrader.feeds as btfeeds
import talib
import numpy as np
from tqdm import tqdm
from my_plot import *
from py_mysql import *

def test_all(code_list):
    objData5 = {
        'x':code_list[0],
        'y':code_list[1],
        'deal_size':0,  #成交次数
        'end_funds':None,  #结束资金
    }
    # 数据查询
    def selectSotck(adjustflag, code, startDate, endDate):
        client = MongoClient('47.100.19.231', 27018)
        db = client.admin
        db.authenticate("quant", "Qweasd123")
        db = client.quant
        mycol = db.testKline

        myquery = {
            "adjustflag": adjustflag,  # 复权类型
            "code": code,  # 股票代码
            "date": {'$gt': startDate, '$lt': endDate}  # 时间区间
        }

        mydoc = mycol.find(myquery)
        data_list = []
        for x in mydoc:
            data_list.append(x)
        fields = ['date', 'open', 'close', 'high', 'low', 'volume', 'code']
        df = pd.DataFrame(data_list, columns=fields)
        test_list = ['open', 'close', 'high', 'low', 'volume']
        df[test_list] = df[test_list].astype('float64')
        df.index = pd.to_datetime(df.date)
        df = df[['open', 'high', 'low', 'close', 'volume']]

        return df

    # 配对股票
    code_list = code_list
    # 定义配对形成期
    formStart = '2019-01-01'
    formEnd = '2019-12-31'
    # 设定交易期时间，选取交易期数据，寻找配对交易开仓和平仓点位(2.1)
    tradStart = '20200101'
    tradEnd = '20211101'

    objData = {}
    for item in code_list:
        df = selectSotck('3', item, '20190101', '20191231')
        objData[item] = df


    def SSD_Spread(priceX, priceY):
        priceX = np.log10(priceX)
        priceY = np.log10(priceY)
        retx = priceX.diff()[1:]
        rety = priceY.diff()[1:]
        standardX = (1 + retx).cumprod()
        standardY = (1 + rety).cumprod()
        spread = standardY - standardX
        return(spread)


    df2 = SSD_Spread(objData[code_list[0]][formStart:formEnd],
                    objData[code_list[1]][formStart:formEnd])


    mu = np.mean(df2)
    sd = np.std(df2)
    mu_close = mu.close
    sd_close = sd.close

    # import matplotlib.pyplot as plt
    # plt.figure(figsize=(12,8))
    # objData2 = {}
    # for item2 in code_list:
    #     df = selectSotck('3',item2,tradStart,tradEnd)
    #     objData2[item2] = df
    # dfff = SSD_Spread(objData2[code_list[0]],objData2[code_list[1]])
    # df3 = dfff.close
    # df3.plot()
    # plt.title('协整配对',loc='center', fontsize=16)
    # plt.axhline(y=mu_close,color='black')
    # plt.axhline(y=mu_close+0.2*sd_close,color='blue',ls='-',lw=2)
    # plt.axhline(y=mu_close-0.2*sd_close,color='blue',ls='-',lw=2)
    # plt.axhline(y=mu_close+2.0*sd_close,color='green',ls='--',lw=2.5)
    # plt.axhline(y=mu_close-2.0*sd_close,color='green',ls='--',lw=2.5)
    # plt.axhline(y=mu_close+3.5*sd_close,color='red',ls='-.',lw=3)
    # plt.axhline(y=mu_close-3.5*sd_close,color='red',ls='-.',lw=3)
    # plt.show()

    # 提前算出信号点，直接可以取
    objData3 = {}
    for item3 in code_list:
        df = selectSotck('3',item3,tradStart,tradEnd)
        objData3[item3] = df
    df_level = SSD_Spread(objData3[code_list[0]],objData3[code_list[1]])
    # 创建策略继承bt.Strategy
    class TestStrategy(bt.Strategy):
        params = (
            # 均线参数设置15天，15日均线
            ('maperiod', 15),
        )

        def log(self, txt, dt=None):
            # 记录策略的执行日志
            dt = dt or self.datas[0].datetime.date(0)
            print('%s, %s' % (dt.isoformat(), txt))

        def __init__(self):
            # 保存收盘价的引用
            self.X_price = self.datas[0].close
            self.Y_price = self.datas[1].close
            # 跟踪挂单
            self.order = None
            # 买入价格和手续费
            self.buyprice = None
            self.buycomm = None

        def next(self):
            self.current_date = self.datas[0].datetime.date(0).isoformat()  # 当前日期
            current_capital = self.broker.getvalue()  # 当前剩余资金

            self.X_size = self.getposition(self.datas[0]).size
            self.Y_size = self.getposition(self.datas[1]).size

            try:
                self.mu_deal = df_level.loc[self.current_date,'close']
            except:
                return

            # print('日期:{}----mu:{}'.format(self.current_date,self.mu_deal))

            if self.mu_deal > (mu_close + 1.2*sd_close):
                if(self.Y_size > -100):
                    self.sell(data=self.datas[1], size=100)

                if(self.X_size < 100):
                    self.buy(data=self.datas[0], size=100)

            if self.mu_deal < (mu_close - 1.2*sd_close):
                if(self.Y_size < 100):
                    self.buy(data=self.datas[1], size=100)

                if(self.X_size > -100):
                    self.sell(data=self.datas[0], size=100)

        # 订单状态通知，买入卖出都是下单

        def notify_order(self, order):
            # 提交了/接受了,  买/卖订单什么都不做
            if order.status in [order.Submitted, order.Accepted]:
                return
            # 检查一个订单是否完成
            # 注意:当资金不足时，broker会拒绝订单
            if order.status in [order.Completed]:
                objData5['deal_size'] += 1
                if order.isbuy():
                    self.buyprice = order.executed.price
                    self.buycomm = order.executed.comm
                    print({'状态': '买入', 'date': self.current_date, 'value': round(
                        order.executed.price, 4), 'size': order.executed.size, 'comm': round(order.executed.comm, 4)})
                elif order.issell():
                    pass
                    print({'状态': '卖出', 'date': self.current_date, 'value': round(
                        order.executed.price, 4), 'size': order.executed.size, 'comm': round(order.executed.comm, 4)})
                self.bar_executed = len(self)

            elif order.status in [order.Canceled, order.Margin, order.Rejected]:
                print('订单取消/保证金不足/拒绝')
            # 其他状态记录为：无法挂单
            self.order = None

        # 交易状态通知，一买一卖算交易（交易净利润）

        def notify_trade(self, trade):
            if not trade.isclosed:
                return


    def main(adjustflag, codeArr, startDate, endDate, startcash, com, qts=100):

        # 创建主控制器
        cerebro = bt.Cerebro()
        # 导入策略参数寻优
        cerebro.addstrategy(TestStrategy)

        # 获取数据
        for item in codeArr:
            df = selectSotck(adjustflag, item, startDate, endDate)
            # 将数据加载至回测系统
            data = bt.feeds.PandasData(dataname=df)
            cerebro.adddata(data)

        # broker设置资金、手续费
        cerebro.broker.setcash(startcash)
        cerebro.broker.setcommission(commission=com)
        # 设置买入设置，策略，数量
        cerebro.addsizer(bt.sizers.FixedSize, stake=qts)
        # 以发出信号当日收盘价成交
        cerebro.broker.set_coc(True)

        print('期初总资金: %.2f' % cerebro.broker.getvalue())
        cerebro.addanalyzer(bt.analyzers.TimeReturn, _name='_TimeReturn')
        result = cerebro.run()
        print('期末总资金: %.2f' % cerebro.broker.getvalue())
        objData5['end_funds'] = cerebro.broker.getvalue()

        # custom_plot(result)
        # cerebro.plot()


    run = main(
        adjustflag='1',  # 复权类型 3：不复权 1：后复权；2：前复权
        codeArr=code_list,  # 股票池
        startDate=tradStart,  # 开始时间
        endDate=tradEnd,  # 结束时间
        startcash=10000000,  # 资金
        com=0.0015  # 手续费
    )

    return objData5

code_all_arr = ['sz.002807','sh.601398','sh.601077','sh.600928','sh.601818','sh.601860','sh.600016','sz.002966','sh.601916','sz.002948']

csv_data = []
for i in tqdm(code_all_arr):
    for j in tqdm(code_all_arr):
        if i == j:
            continue
        else:
            test = test_all([i,j])
            csv_data.append(test)

print(csv_data)
# (pd.DataFrame(csv_data)).to_csv('test.csv',sep=',',index=True,header=True)